package main

import (
	pb "03-grpc-stream/proto"
	"fmt"
	"net"
	"sync"
	"time"

	"google.golang.org/grpc"
)

const PORT = ":3001"

type server struct {
	pb.UnimplementedStreamServer
}

func (s *server) GetStream(req *pb.ReqMessage, res grpc.ServerStreamingServer[pb.ResMessage]) error {
	fmt.Println(req.Data)
	i := 0
	for {
		i++
		_ = res.Send(&pb.ResMessage{
			Data: fmt.Sprintf("%v", time.Now().Format("2006-01-02 15:04:05")),
		})
		time.Sleep(time.Second)
		if i > 5 {
			break
		}
	}
	return nil
}

func (s *server) PutStream(clientStr grpc.ClientStreamingServer[pb.ReqMessage, pb.ResMessage]) error {
	for {
		if a, err := clientStr.Recv(); err != nil {
			fmt.Println(err)
			break
		} else {
			fmt.Println(a.Data)
		}
	}
	return nil
}

func (s *server) AllStream(allStr grpc.BidiStreamingServer[pb.ReqMessage, pb.ResMessage]) error {
	wg := sync.WaitGroup{}
	wg.Add(2)
	go func() {
		defer wg.Done()
		for {
			data, _ := allStr.Recv()
			fmt.Println("收到客户端数据: " + data.Data)
		}
	}()

	go func() {
		defer wg.Done()
		for {
			_ = allStr.Send(&pb.ResMessage{Data: "server message"})
			time.Sleep(2 * time.Second)
		}
	}()

	wg.Wait()
	return nil
}

func main() {
	lis, err := net.Listen("tcp", PORT)
	if err != nil {
		panic(err)
	}
	s := grpc.NewServer()
	pb.RegisterStreamServer(s, &server{})
	err = s.Serve(lis)
	if err != nil {
		panic(err)
	}
}
